#!/usr/bin/env python3
"""将 data 中的 PDF 转换为 Markdown，使用 markitdown 包。

生成文件：data/2016年云南农民工调查报告.md
"""
from pathlib import Path
import sys
import argparse
import os

ROOT = Path(__file__).resolve().parents[1]
PDF_PATH = ROOT / 'data' / '2016年云南农民工调查报告.pdf'
OUT_MD = ROOT / 'data' / '2016年云南农民工调查报告.md'


def convert_with_markitdown(pdf_path: Path, out_path: Path, enable_ocr: bool = False, extract_images: bool = False, use_llm: bool = False) -> bool:
    try:
        import markitdown
    except Exception as e:
        print('markitdown not available:', e)
        return False
    llm_client = None
    if use_llm:
        api_key = os.environ.get('OPENAI_API_KEY')
        if api_key:
            try:
                import openai
                openai.api_key = api_key
                llm_client = openai
                print('Initialized openai client from OPENAI_API_KEY')
            except Exception as e:
                print('Failed to init openai client:', e)
                llm_client = None
        else:
            print('OPENAI_API_KEY not found in environment; skipping llm client')

    try:
        md = markitdown.MarkItDown()
        kwargs = {}
        if enable_ocr:
            kwargs['enable_ocr'] = True
        if extract_images:
            kwargs['extract_images'] = True
        if llm_client is not None:
            kwargs['llm_client'] = llm_client

        res = md.convert(str(pdf_path), **kwargs)
        markdown_text = getattr(res, 'markdown', None)
        if markdown_text is None:
            markdown_text = str(res)
        out_path.write_text(markdown_text, encoding='utf-8')
        print('Saved', out_path)
        return True
    except Exception as e:
        print('markitdown conversion failed:', repr(e))
        return False


def fallback_extract_text(pdf_path: Path, out_path: Path) -> bool:
    """简单回退：使用 PyPDF2 提取纯文本并写入 markdown 文件（如果可用）。"""
    try:
        import PyPDF2
    except Exception as e:
        print('PyPDF2 not available for fallback:', e)
        return False

    try:
        text_parts = []
        with pdf_path.open('rb') as f:
            reader = PyPDF2.PdfReader(f)
            for p in reader.pages:
                try:
                    txt = p.extract_text() or ''
                except Exception:
                    txt = ''
                text_parts.append(txt)
        content = '\n\n'.join(part.strip() for part in text_parts if part and part.strip())
        if not content:
            print('Fallback extraction produced no text')
            return False
        out_path.write_text(content, encoding='utf-8')
        print('Saved (fallback) ', out_path)
        return True
    except Exception as e:
        print('Fallback extraction failed:', repr(e))
        return False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--enable-ocr', action='store_true', help='Enable OCR during conversion')
    parser.add_argument('--extract-images', action='store_true', help='Tell markitdown to extract images')
    parser.add_argument('--use-llm', action='store_true', help='Attempt to initialize an LLM client from OPENAI_API_KEY and pass to markitdown')
    args = parser.parse_args()

    if not PDF_PATH.exists():
        print('PDF not found:', PDF_PATH)
        sys.exit(2)

    OK = convert_with_markitdown(PDF_PATH, OUT_MD, enable_ocr=args.enable_ocr, extract_images=args.extract_images, use_llm=args.use_llm)
    if not OK:
        print('Trying fallback extraction...')
        OK = fallback_extract_text(PDF_PATH, OUT_MD)

    if not OK:
        print('Conversion failed')
        sys.exit(1)


if __name__ == '__main__':
    main()
